package de.dopichaj.labrador.util.io.huffman;

import java.io.IOException;
import java.math.BigInteger;
import java.util.HashMap;
import java.util.Map;
import java.util.PriorityQueue;

import de.dopichaj.labrador.util.io.BitInputStream;
import de.dopichaj.labrador.util.io.BitOutputStream;


public final class HuffmanTree<T> {
    
    static final class HuffmanCode {
       
        /**
         * The code bits (only the lower <code>significantBits</code> are relevant).
         */
        final BigInteger code;
        final int significantBits;
        
        public HuffmanCode(final BigInteger code, final int significantBits) {
            
            assert significantBits > 0
                : "Invalid number of significant bits: " + significantBits;
            this.code = code;
            this.significantBits = significantBits;
        }
    }

    /**
     * A map from the symbols to the corresponding Huffman codes.
     * The codes are guaranteed to be suffix-free.
     */
    private final Map<T, HuffmanCode> map = new HashMap<T, HuffmanCode>();
    private final HuffmanNode<T> root;
    
    /**
     * 
     * @param countMap a map from the symbols to their frequency
     */
    public HuffmanTree(final Map<T, Integer> countMap) {
        
        // if inserting more than one symbol
        if (countMap.size() > 1) {
            
            // construct a heap from the leaf nodes
            final PriorityQueue<HuffmanNode<T>> heap =
                new PriorityQueue<HuffmanNode<T>>(countMap.size());
            for (final Map.Entry<T, Integer> entry : countMap.entrySet()) {
                heap.add(new LeafHuffmanNode<T>(entry.getKey(), entry.getValue()));
            }
            
            // while the heap has more than one item
            while (heap.size() > 1) {
                
                // remove the first two items and insert a new inner node
                heap.add(new InnerHuffmanNode<T>(heap.poll(), heap.poll()));
            }
            
            // the single remaining node is the root of the Huffman tree
            root = heap.poll();
            
            // traverse the tree and build a mapping from value to Huffman code
            root.addToCodeMap(map, BigInteger.ZERO, 0);
        
        // else if there is exactly one symbol
        } else if (countMap.size() == 1) {
            
            final LeafHuffmanNode<T> leafNode =
                new LeafHuffmanNode<T>(countMap.keySet().iterator().next(), 1);
            // use the single bit 0 to encode the symbol
            map.put(leafNode.getData(), new HuffmanCode(BigInteger.ZERO, 1));
            
            // build a dummy tree
            root = new InnerHuffmanNode<T>(leafNode, leafNode);
            
        // else there are no symbols
        } else {
            
            // the map is empty
            
            // the tree is empty
            root = null;
        }
    }
    
    /**
     * Create a Huffman tree from the input, as written by {@link #writeTree(BitOutputStream, HuffmanDataWriter)}.
     * The resulting tree is read-only.
     */
    public HuffmanTree(final BitInputStream input, final HuffmanDataReader<T> reader) throws IOException {
        root = readHuffmanNode(input, reader);
    }

    private HuffmanNode<T> readHuffmanNode(BitInputStream input, HuffmanDataReader<T> reader) throws IOException {
        
        final int type = input.readBits(1);
        if (type == 0) {
            
            final HuffmanNode<T> child0 = readHuffmanNode(input, reader);
            final HuffmanNode<T> child1 = readHuffmanNode(input, reader);
            return new InnerHuffmanNode<T>(child0, child1);
            
        } else {
            
            final T symbol = reader.read(input);
            return new LeafHuffmanNode<T>(symbol, 0);
        }
    }
    
    /**
     * Write the Huffman code corresponding to the given symbol to the output
     * stream.
     * 
     * @param symbol the symbol to encode; it must be in the tree
     * @param output the stream to write to
     * @throws IOException if an I/O error occurs
     */
    public void writeCode(final T symbol, final BitOutputStream output) throws IOException {
        
        final HuffmanCode code = map.get(symbol);
        output.writeBits(code.code, code.significantBits);
    }
    
    /**
     * Read the symbol that corresponds to the following bits in the input stream.
     * @throws IOException 
     */
    public T readSymbol(final BitInputStream input) throws IOException {

        return root.read(input);
    }
    
    /**
     * Write the encoding tree to the output.
     */
    public void writeTree(final BitOutputStream output, final HuffmanDataWriter<T> writer) throws IOException {
        root.write(output, writer);
        output.close();
    }

    /**
     * Generate a string representation of the tree's encodings.
     * Warning: this is an expensive operation and yields a multi-line string.
     */
    @Override
    public String toString() {
        final StringBuilder builder = new StringBuilder();
        for (final Map.Entry<T, HuffmanCode> entry : map.entrySet()) {
            builder.append(entry.getKey() + "\t");
            final HuffmanCode code = entry.getValue();
            BigInteger current = code.code;
            for (int i = 0; i < code.significantBits; i++) {
                builder.append(current.and(BigInteger.ONE));
                current = current.shiftRight(1);
            }
            builder.append("\n");
        }
        return builder.toString();
    }
}
/*
Copyright (c) 2007 Philipp Dopichaj

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
*/